package kafka

import (
	"context"

	"gitee.com/zawei/yggdrasil/pkg/log"
	"github.com/Shopify/sarama"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
)

type Message struct {
	*sarama.ConsumerMessage
	Metadata string
}

type MessageHandle func(ctx context.Context, msg *Message)

func traceConsumerWrapper(handle MessageHandle) MessageHandle {
	return func(ctx context.Context, msg *Message) {
		headers := make(map[string]string)
		for _, header := range msg.Headers {
			headers[string(header.Key)] = string(header.Value)
		}
		key, topic, partition, offset := msg.Key, msg.Topic, msg.Partition, msg.Offset
		spanContext, _ := opentracing.GlobalTracer().Extract(
			opentracing.TextMap,
			opentracing.TextMapCarrier(headers),
		)
		span, ctx := opentracing.StartSpanFromContext(ctx,
			"kafka.group_consumer."+topic,
			opentracing.FollowsFrom(spanContext),
			ext.SpanKindConsumer)
		ext.Component.Set(span, "go-kafka")
		ext.PeerService.Set(span, "kafka")
		span.SetTag("topic", topic)
		span.SetTag("partition", partition)
		span.SetTag("key", string(key))
		span.SetTag("offset", offset)
		defer span.Finish()
		handle(ctx, msg)
	}
}

func offsetConsumerWrapper(offsetManager OffsetManager, sess sarama.ConsumerGroupSession, handle MessageHandle) MessageHandle {
	return func(ctx context.Context, msg *Message) {
		offsetManager.Read(msg)
		defer offsetManager.Commit(sess, msg)
		handle(ctx, msg)
	}
}

type MessageHandler interface {
	Setup(sess sarama.ConsumerGroupSession) error
	Cleanup(sess sarama.ConsumerGroupSession) error
	HandleMsg(ctx context.Context, msg *Message)
}

type messageHandler struct{}

func (messageHandler) Setup(sarama.ConsumerGroupSession) error   { return nil }
func (messageHandler) Cleanup(sarama.ConsumerGroupSession) error { return nil }
func (messageHandler) HandleMsg(context.Context, *Message)       {}

type Consumer struct {
	*GroupConsumerConfig

	group         sarama.ConsumerGroup
	offsetManager OffsetManager
}

func (s *Consumer) Consume(ctx context.Context) {
	for {
		if err := s.group.Consume(ctx, s.Topics, s); err != nil {
			log.Errorw("kafka consumer", "error", err)
		}
		if ctx.Err() != nil {
			return
		}
	}

}

func (s *Consumer) Stop() error {
	return s.group.Close()
}

func (s *Consumer) Setup(sess sarama.ConsumerGroupSession) error {
	log.Infow("consumerGroupHandler setup", "claims", sess.Claims())
	return s.handler.Setup(sess)
}

func (s *Consumer) Cleanup(sess sarama.ConsumerGroupSession) error {
	log.Infow("consumerGroupHandler cleanup", "claims", sess.Claims())
	return s.handler.Cleanup(sess)
}

func (s *Consumer) ConsumeClaim(sess sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	for {
		msg, ok := <-claim.Messages()
		if !ok {
			return nil
		}
		log.Debugw("consume claim", "topic", msg.Topic, "partition", msg.Partition, "offset", msg.Offset)
		var handle = offsetConsumerWrapper(s.offsetManager, sess, s.handler.HandleMsg)
		if s.EnableTrace {
			handle = traceConsumerWrapper(handle)
		}
		handle(context.TODO(), &Message{ConsumerMessage: msg})
	}
}
